from parse import stmt, expr
from typing import List
from .token import Token, TokenType


class Parser:
    """
    program ::= declaration* EOF;
    declaration ::= classDecl
                    | funDecl
                    | varDeclaration
                    | statement;
    classDecl = "class" IDENTIFIER ("<" IDENTIFIER)? "{" function* "}"
    funDecl ::= "fun" function;
    function = IDENTIFIER "(" parameters? ")" block ;
    varDeclaration ::= "var" IDENTIFIER ( "=" expression )? ";";
    statement ::= exprStmt
                  | forStmt
                  | returnStmt
                  | breakStmt
                  | continueStmt
                  | ifStmt
                  | printStmt
                  | whileStmt
                  | block;
    ifStmt = "if" "(" expression ")" statement ("else" statement )? ;
    exprStmt ::= expression ";";
    forStmt ::= "for" "("(varDecl | exprStmt ";") expression? ; expression?; ")" statement ;
    returnStmt ::= "return" expression? ";";
    printStmt ::= "print" expression ";";
    whileStmt ::= "while" "(" expression ")" statement;
    block ::= "{" declaration* "}";
    expression ::= assignment;
    assignment ::= ( call ".")? IDENTIFIER "=" assignment | logic_or;
    logic_or ::= logic_and ( "or" logic_and)* ;
    logic_and ::= equality ( "and" equality)* ;
    equality ::= comparison ( ( "!=" | "==" ) comparison )* ;
    comparison ::= term ( ( ">" | ">=" | "<" | "<=" ) term )* ;
    term ::= factor ( ( "-" | "+" ) factor )* ;
    factor ::= unary ( ( "/" | "*" ) unary )* ;
    unary ::= ( "!" | "-" ) unary | call ;
    call ::= primary ( "(" arguments? ")" | "." IDENTIFIER )*
    arguments ::= expression ("," expression )*
    parameters ::= IDENTIFIER ( "," IDENTIFIER )* ;
    primary ::= NUMBER
                | STRING
                | "true"
                | "false"
                | "nil"
                | "this"
                | "(" expression ")"
                | IDENTIFIER
                | "super" "." IDENTIFIER;
    """

    def __init__(self, tokens: List[Token]):
        self.tokens = tokens
        self.current = 0

    def parse(self):
        stmts = []
        while not self.is_eof():
            stmts.append(self.declaration())

        return stmts

    def declaration(self):
        if self.match_token(TokenType.CLASS):
            return self.class_declaration()
        elif self.match_token(TokenType.FUN):
            return self.function(stmt.FunType.FUN)
        elif self.match_token(TokenType.VAR):
            return self.var_declaration()
        else:
            return self.statement()

    def class_declaration(self):
        name = self.consume(TokenType.IDENTIFIER)
        if self.match_token(TokenType.LESS):
            superclass = self.consume(TokenType.IDENTIFIER)
        else:
            superclass = None

        self.consume(TokenType.LEFT_BRACE)

        methods = []

        while not self.check(TokenType.RIGHT_BRACE):
            methods.append(self.function(stmt.FunType.METHOD))

        self.consume(TokenType.RIGHT_BRACE)

        return stmt.StmtClass(name, superclass, methods)

    def var_declaration(self):
        name = self.consume(TokenType.IDENTIFIER)
        if self.match_token(TokenType.EQUAL):
            initializer = self.expression()
        else:
            initializer = expr.ExprLiteral(None)

        self.consume(TokenType.SEMICOLON)

        return stmt.StmtVar(name, initializer)

    def function(self, kind: stmt.FunType):
        name = self.consume(TokenType.IDENTIFIER)
        self.consume(TokenType.LEFT_PARENT)

        params = []
        if self.check(TokenType.IDENTIFIER):
            params.append(self.consume(TokenType.IDENTIFIER))

        while self.check(TokenType.COMMA):
            self.consume(TokenType.COMMA)
            params.append(self.consume(TokenType.IDENTIFIER))

        self.consume(TokenType.RIGHT_PARENT)

        self.consume(TokenType.LEFT_BRACE)
        body = []
        while not self.check(TokenType.RIGHT_BRACE):
            body.append(self.declaration())

        # 函数末尾默认返回Nil
        body.append(stmt.StmtReturn(expr.ExprLiteral(None)))

        if self.is_eof():
            raise ValueError(f'unclosed function')

        self.consume(TokenType.RIGHT_BRACE)

        return stmt.StmtFun(name, params, body, kind)

    def statement(self):
        if self.match_token(TokenType.FOR):
            return self.for_stmt()
        elif self.match_token(TokenType.RETURN):
            return self.return_stmt()
        elif self.match_token(TokenType.CONTINUE):
            return self.continue_stmt()
        elif self.match_token(TokenType.BREAK):
            return self.break_stmt()
        elif self.match_token(TokenType.PRINT):
            return self.print_stmt()
        elif self.match_token(TokenType.WHILE):
            return self.while_stmt()
        elif self.match_token(TokenType.IF):
            return self.if_stmt()
        else:
            return self.expr_stmt()

    def for_stmt(self):
        self.consume(TokenType.LEFT_PARENT)
        # 初始化
        if self.match_token(TokenType.SEMICOLON):
            init_stmt = None
        elif self.match_token(TokenType.VAR):
            init_stmt = self.var_declaration()
        else:
            init_stmt = self.statement()

        # 条件
        if self.match_token(TokenType.SEMICOLON):
            condition = expr.ExprLiteral(True)
        else:
            condition = self.expression()
            self.consume(TokenType.SEMICOLON)

        # 自增语句
        if self.match_token(TokenType.RIGHT_PARENT):
            inc_stmt = None
        else:
            inc_stmt = stmt.StmtExpr(self.expression())
            self.consume(TokenType.RIGHT_PARENT)

        blocks = []
        if init_stmt:
            blocks.append(init_stmt)

        self.consume(TokenType.LEFT_BRACE)

        while not self.check(TokenType.RIGHT_BRACE) and not self.is_eof():
            blocks.append(self.statement())

        self.consume(TokenType.RIGHT_BRACE)

        if inc_stmt:
            blocks.append(inc_stmt)

        while_stmt = stmt.StmtWhile(condition, stmt.StmtBlock(blocks))

        if init_stmt:
            return stmt.StmtBlock([init_stmt, while_stmt])
        else:
            return while_stmt

    def consume(self, t: TokenType):
        if self.peek().token_type != t:
            raise ValueError(f'expect {t=}, but {self.peek()}')

        return self.advance()

    def peek(self):
        return self.tokens[self.current]

    def match_token(self, t: TokenType):
        if self.is_eof():
            return False

        if self.peek().token_type == t:
            self.advance()
            return True

        return False

    def match_tokens(self, *types: TokenType):
        for i, t in enumerate(types):
            if self.tokens[self.current + i].token_type != t:
                return False

        self.current += len(types)

        return True

    def advance(self):
        ret = self.peek()
        self.current += 1

        return ret

    def check(self, *t: TokenType):
        for i, v in enumerate(t):
            if self.tokens[self.current + i].token_type != v:
                return False

        return True

    def is_eof(self):
        return self.peek().token_type == TokenType.EOF

    def expression(self):
        return self.assignment()

    def return_stmt(self):
        return_expr = None if self.match_token(TokenType.SEMICOLON) else self.expression()
        self.consume(TokenType.SEMICOLON)
        return stmt.StmtReturn(return_expr)

    def print_stmt(self):
        ret = stmt.StmtPrint(self.expression())
        self.consume(TokenType.SEMICOLON)
        return ret

    def while_stmt(self):
        self.consume(TokenType.LEFT_PARENT)
        condition = self.expression()
        self.consume(TokenType.RIGHT_PARENT)

        blocks = []

        self.consume(TokenType.LEFT_BRACE)
        while not self.check(TokenType.RIGHT_BRACE) and not self.is_eof():
            blocks.append(self.statement())

        self.consume(TokenType.RIGHT_BRACE)

        return stmt.StmtWhile(condition, stmt.StmtBlock(blocks))

    def expr_stmt(self):
        ret = stmt.StmtExpr(self.expression())
        self.consume(TokenType.SEMICOLON)
        return ret

    def if_stmt(self):
        self.consume(TokenType.LEFT_PARENT)
        condition = self.expression()
        self.consume(TokenType.RIGHT_PARENT)
        self.consume(TokenType.LEFT_BRACE)

        then_block = []

        while not self.check(TokenType.RIGHT_BRACE) and not self.is_eof():
            then_block.append(self.statement())

        self.consume(TokenType.RIGHT_BRACE)

        else_block = []

        if self.match_token(TokenType.ELSE):
            self.consume(TokenType.LEFT_BRACE)
            while not self.check(TokenType.RIGHT_BRACE) and not self.is_eof():
                else_block.append(self.statement())

            self.consume(TokenType.RIGHT_BRACE)

        if_block = stmt.StmtBlock(then_block)
        else_block = stmt.StmtBlock(else_block) if else_block else None

        return stmt.StmtIf(condition, if_block, else_block)

    def assignment(self) -> expr.Expr:
        or_expr = self.logic_or()
        if self.match_token(TokenType.EQUAL):
            value = self.assignment()
            if isinstance(or_expr, expr.ExprVariable):
                return expr.ExprAssign(or_expr.name, value)
            elif isinstance(or_expr, expr.ExprGet):
                return expr.ExprSet(or_expr.obj, or_expr.name, value)
            else:
                raise TypeError(f'expected expression in assignment {or_expr}')
        else:
            return or_expr

    def logic_or(self) -> expr.Expr:
        ret = self.logic_and()

        while not self.is_eof() and self.check(TokenType.OR):
            ret = expr.ExprLogical(ret, self.advance(), self.logic_and())

        return ret

    def logic_and(self) -> expr.Expr:
        ret = self.equality()

        while not self.is_eof() and self.check(TokenType.AND):
            ret = expr.ExprLogical(ret, self.advance(), self.equality())

        return ret

    def equality(self) -> expr.Expr:
        ret = self.comparison()

        while not self.is_eof() and (self.check(TokenType.BANG_EQUAL) or self.check(TokenType.EQUAL_EQUAL)):
            ret = expr.ExprLogical(ret, self.advance(), self.comparison())

        return ret

    def comparison(self) -> expr.Expr:
        ret = self.term()

        while not self.is_eof() \
                and (self.check(TokenType.GREATER)
                     or self.check(TokenType.GREATER_EQUAL)
                     or self.check(TokenType.LESS)
                     or self.check(TokenType.LESS_EQUAL)):
            ret = expr.ExprBinary(ret, self.advance(), self.term())

        return ret

    def term(self) -> expr.Expr:
        ret = self.factor()

        while not self.is_eof() and (self.check(TokenType.PLUS) or self.check(TokenType.MINUS)):
            ret = expr.ExprBinary(ret, self.advance(), self.factor())

        return ret

    def factor(self) -> expr.Expr:
        ret = self.unary()

        while not self.is_eof() and (self.check(TokenType.PLUS) or self.check(TokenType.MINUS)):
            ret = expr.ExprBinary(ret, self.advance(), self.unary())

        return ret

    def unary(self) -> expr.Expr:
        if self.check(TokenType.BANG) or self.check(TokenType.MINUS):
            return expr.ExprUnary(self.advance(), self.unary())
        else:
            return self.call()

    def call(self) -> expr.Expr:
        ret = self.primary()

        while True:
            if self.match_token(TokenType.LEFT_PARENT):
                ret = self.finish_call(ret)
            elif self.match_token(TokenType.DOT):
                ret = expr.ExprGet(ret, self.consume(TokenType.IDENTIFIER))
            else:
                break

        return ret

    def finish_call(self, e: expr.Expr) -> expr.Expr:
        args = []
        if not self.match_token(TokenType.RIGHT_PARENT):
            args.append(self.expression())

            while not self.is_eof() and not self.check(TokenType.RIGHT_PARENT):
                self.consume(TokenType.COMMA)
                args.append(self.expression())

            self.consume(TokenType.RIGHT_PARENT)

        return expr.ExprCall(e, args)

    def primary(self) -> expr.Expr:
        if self.check(TokenType.NUMBER):
            return expr.ExprLiteral(float(self.advance().lexeme))
        elif self.check(TokenType.STRING):
            return expr.ExprLiteral(self.advance().lexeme)
        elif self.match_token(TokenType.TRUE):
            return expr.ExprLiteral(True)
        elif self.match_token(TokenType.FALSE):
            return expr.ExprLiteral(False)
        elif self.match_token(TokenType.NIL):
            return expr.ExprLiteral(None)
        elif self.check(TokenType.THIS):
            return expr.This(self.advance())
        elif self.match_token(TokenType.LEFT_PARENT):
            ret = self.expression()
            self.consume(TokenType.RIGHT_PARENT)
            return ret
        elif self.check(TokenType.IDENTIFIER):
            return expr.ExprVariable(self.advance())
        elif self.check(TokenType.SUPER, TokenType.DOT, TokenType.IDENTIFIER):
            self.advance()
            self.advance()
            return expr.Super(self.advance())
        else:
            raise ValueError(f'unexpected primary token {self.advance()}')

    def continue_stmt(self):
        ret = stmt.StmtContinue()
        self.consume(TokenType.SEMICOLON)
        return ret

    def break_stmt(self):
        ret = stmt.StmtBreak()
        self.consume(TokenType.SEMICOLON)
        return ret
